#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PostgreSQL数据库迁移脚本

将现有的SQLite数据库迁移到PostgreSQL，或者在PostgreSQL中创建新的数据库表结构。

使用方法:
    python migrate_to_postgresql.py --create-only  # 仅创建表结构
    python migrate_to_postgresql.py --migrate      # 迁移数据
    python migrate_to_postgresql.py --reset        # 重置数据库

@author: wechat_robot
@created: 2025-08-28
"""

import os
import sys
import argparse
import logging
from pathlib import Path
from datetime import datetime

# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

try:
    import psycopg2
    from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
except ImportError:
    print("错误: 未安装 psycopg2-binary")
    print("请运行: pip install psycopg2-binary")
    sys.exit(1)

from api_restructured import create_app
from models.base import db
from models import User, Message, Group, Config, WelcomeRecord, Belong
from utils.database.config import DatabaseConfig
from utils.database.database import init_database, create_tables, init_default_configs

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def get_postgresql_config():
    """
    获取PostgreSQL配置
    
    Returns:
        dict: PostgreSQL配置字典
    """
    print("=== PostgreSQL数据库配置 ===")
    
    # 尝试从环境变量获取配置，默认使用测试环境配置
    config = {
        'type': 'postgresql',
        'host': os.getenv('DB_HOST', '192.168.0.141'),
        'port': int(os.getenv('DB_PORT', 15432)),
        'database': os.getenv('DB_NAME', 'aisales'),
        'username': os.getenv('DB_USER', 'chatwiki'),
        'password': os.getenv('DB_PASSWORD', 'postgres_password'),
        'options': {
            'sslmode': 'disable',
            'connect_timeout': 10,
            'application_name': 'wechat_robot'
        }
    }
    
    # 如果环境变量中没有密码，提示用户输入
    if not config['password']:
        import getpass
        config['password'] = getpass.getpass(f"请输入PostgreSQL密码 (用户: {config['username']}): ")
    
    return config


def test_postgresql_connection(config):
    """
    测试PostgreSQL连接
    
    Args:
        config (dict): 数据库配置
        
    Returns:
        bool: 连接是否成功
    """
    try:
        conn = psycopg2.connect(
            host=config['host'],
            port=config['port'],
            database=config['database'],
            user=config['username'],
            password=config['password']
        )
        conn.close()
        logger.info("✓ PostgreSQL连接测试成功")
        return True
    except psycopg2.OperationalError as e:
        logger.error(f"✗ PostgreSQL连接失败: {e}")
        return False
    except Exception as e:
        logger.error(f"✗ 连接测试出错: {e}")
        return False


def create_database_if_not_exists(config):
    """
    如果数据库不存在则创建
    
    Args:
        config (dict): 数据库配置
    """
    try:
        # 连接到postgres数据库（默认存在）
        conn = psycopg2.connect(
            host=config['host'],
            port=config['port'],
            database='postgres',
            user=config['username'],
            password=config['password']
        )
        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        cursor = conn.cursor()
        
        # 检查数据库是否存在
        cursor.execute(
            "SELECT 1 FROM pg_catalog.pg_database WHERE datname = %s",
            (config['database'],)
        )
        
        if cursor.fetchone() is None:
            # 数据库不存在，创建它
            cursor.execute(f'CREATE DATABASE "{config["database"]}"')
            logger.info(f"✓ 数据库 '{config['database']}' 创建成功")
        else:
            logger.info(f"✓ 数据库 '{config['database']}' 已存在")
        
        cursor.close()
        conn.close()
        
    except Exception as e:
        logger.error(f"✗ 创建数据库失败: {e}")
        raise


def create_tables_postgresql(config):
    """
    在PostgreSQL中创建表结构
    
    Args:
        config (dict): 数据库配置
    """
    logger.info("开始在PostgreSQL中创建表结构...")
    
    # 配置PostgreSQL数据库URL
    database_url = DatabaseConfig.get_database_url(config)
    engine_options = DatabaseConfig.get_engine_options('postgresql')
    
    # 设置环境变量，让应用使用PostgreSQL配置
    os.environ['SQLALCHEMY_DATABASE_URI'] = database_url
    os.environ['DB_TYPE'] = 'postgresql'
    
    # 创建Flask应用（会自动使用环境变量中的数据库配置）
    app = create_app()
    
    with app.app_context():
        try:
            # 创建所有表
            db.create_all()
            logger.info("✓ 数据库表创建成功")
            
            # 初始化默认配置
            init_default_configs()
            logger.info("✓ 默认配置初始化完成")
            
            # 显示创建的表
            inspector = db.inspect(db.engine)
            tables = inspector.get_table_names()
            logger.info(f"✓ 已创建表: {', '.join(tables)}")
            
        except Exception as e:
            logger.error(f"✗ 创建表失败: {e}")
            raise


def reset_postgresql_database(config):
    """
    重置PostgreSQL数据库（删除所有表）
    
    Args:
        config (dict): 数据库配置
    """
    confirm = input("⚠️  警告：此操作将删除PostgreSQL数据库中的所有表和数据，是否继续？(y/N): ")
    if confirm.lower() != 'y':
        logger.info("操作已取消")
        return
    
    logger.info("开始重置PostgreSQL数据库...")
    
    # 配置PostgreSQL数据库URL
    database_url = DatabaseConfig.get_database_url(config)
    engine_options = DatabaseConfig.get_engine_options('postgresql')
    
    # 设置环境变量，让应用使用PostgreSQL配置
    os.environ['SQLALCHEMY_DATABASE_URI'] = database_url
    os.environ['DB_TYPE'] = 'postgresql'
    
    # 创建Flask应用（会自动使用环境变量中的数据库配置）
    app = create_app()
    
    with app.app_context():
        try:
            # 删除所有表
            db.drop_all()
            logger.info("✓ 所有表已删除")
            
            # 重新创建表
            db.create_all()
            logger.info("✓ 表结构重新创建完成")
            
            # 初始化默认配置
            init_default_configs()
            logger.info("✓ 默认配置初始化完成")
            
        except Exception as e:
            logger.error(f"✗ 重置数据库失败: {e}")
            raise


def migrate_data_from_sqlite(config, sqlite_path=None):
    """
    从SQLite迁移数据到PostgreSQL
    
    Args:
        config (dict): PostgreSQL配置
        sqlite_path (str): SQLite数据库文件路径
    """
    if sqlite_path is None:
        sqlite_path = 'wechat_robot.db'
    
    if not os.path.exists(sqlite_path):
        logger.warning(f"SQLite数据库文件不存在: {sqlite_path}")
        logger.info("将仅创建表结构，不进行数据迁移")
        create_tables_postgresql(config)
        return
    
    logger.info(f"开始从SQLite迁移数据到PostgreSQL...")
    logger.info(f"SQLite文件: {sqlite_path}")
    
    # TODO: 实现数据迁移逻辑
    # 这里可以添加具体的数据迁移代码
    logger.warning("数据迁移功能尚未实现，仅创建表结构")
    create_tables_postgresql(config)


def main():
    """
    主函数
    """
    parser = argparse.ArgumentParser(description='PostgreSQL数据库迁移工具')
    parser.add_argument('--create-only', action='store_true', help='仅创建表结构')
    parser.add_argument('--migrate', action='store_true', help='从SQLite迁移数据')
    parser.add_argument('--reset', action='store_true', help='重置数据库')
    parser.add_argument('--sqlite-path', type=str, help='SQLite数据库文件路径')
    
    args = parser.parse_args()
    
    if not any([args.create_only, args.migrate, args.reset]):
        parser.print_help()
        return
    
    try:
        # 获取PostgreSQL配置
        config = get_postgresql_config()
        
        # 测试连接
        if not test_postgresql_connection(config):
            logger.error("PostgreSQL连接失败，请检查配置")
            return
        
        # 确保数据库存在
        create_database_if_not_exists(config)
        
        if args.reset:
            reset_postgresql_database(config)
        elif args.create_only:
            create_tables_postgresql(config)
        elif args.migrate:
            migrate_data_from_sqlite(config, args.sqlite_path)
        
        logger.info("操作完成！")
        
    except KeyboardInterrupt:
        logger.info("操作被用户取消")
    except Exception as e:
        logger.error(f"操作失败: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()